package org.idea.qiyu.framework.apm.agent.core.jdbc;

import javassist.*;

import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.IllegalClassFormatException;
import java.lang.instrument.Instrumentation;
import java.security.ProtectionDomain;
import java.util.HashMap;
import java.util.Map;
import java.util.UUID;

/**
 * druid数据库连接池的agent
 *
 * @Author linhao
 * @Date created in 12:30 下午 2021/7/19
 */
public class DruidPoolAgent {

    private static final String DRUID_POOL_CONN_CLASS = "com/alibaba/druid/pool/DruidPooledConnection";
    private static final String DRUID_POOL_STATEMENT_CLASS = "com/alibaba/druid/pool/DruidPooledStatement";

    private static Map<String, JdbcStatisticsObject> jdbcStatisticsObjectMap = new HashMap<>();

    public static void premain(String args, Instrumentation instrumentation) {
        instrumentation.addTransformer(new ClassFileTransformer() {
            @Override
            public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] classfileBuffer) throws IllegalClassFormatException {
                if (loader == null || className == null) {
                    return null;
                }
                if (DRUID_POOL_CONN_CLASS.equals(className)) {
                    return buildDruidPoolConnectClassByte(className.replace("/", "."));
                }
                if (DRUID_POOL_STATEMENT_CLASS.equals(className)) {
                    return buildDruidPoolStatementClassByte(className.replace("/", "."));
                }
                return null;
            }
        }, true);
    }

    public static byte[] buildDruidPoolStatementClassByte(String className) {
        try {
            ClassPool classPool = ClassPool.getDefault();
            classPool.appendSystemPath();
            CtClass ctClass = classPool.getCtClass(className);

            CtMethod transactionRecordMethod = ctClass.getDeclaredMethod("transactionRecord");
            CtMethod transactionRecordMethodCopyMethod = CtNewMethod.copy(transactionRecordMethod, ctClass, new ClassMap());
            String newBeforeExecuteMethodName = transactionRecordMethod.getName() + "$Agent";
            transactionRecordMethod.setName(newBeforeExecuteMethodName);
            transactionRecordMethodCopyMethod.setBody(
                    "{" +
                            "" + newBeforeExecuteMethodName + "($$);" +
                            "org.idea.qiyu.framework.apm.agent.core.jdbc.DruidPoolAgent.recordSql($$);" +
                            "}");
            ctClass.addMethod(transactionRecordMethodCopyMethod);
            return ctClass.toBytecode();
        } catch (Exception e) {
            e.printStackTrace();
        }
        return null;
    }

    public static void main(String[] args) {
        byte i = 49;
        System.out.println(i);
    }

    public static byte[] buildDruidPoolConnectClassByte(String className) {
        try {
            ClassPool classPool = ClassPool.getDefault();
            classPool.appendSystemPath();
            CtClass ctClass = classPool.getCtClass(className);

            CtMethod beforeExecuteMethod = ctClass.getDeclaredMethod("beforeExecute");
            CtMethod beforeExecuteMethodCopyMethod = CtNewMethod.copy(beforeExecuteMethod, ctClass, new ClassMap());
            String newBeforeExecuteMethodName = beforeExecuteMethod.getName() + "$Agent";
            beforeExecuteMethod.setName(newBeforeExecuteMethodName);
            beforeExecuteMethodCopyMethod.setBody(
                    "{" +
                            "String jdbcUrl = this.holder.dataSource.getRawJdbcUrl();" +
                            "org.idea.qiyu.framework.apm.agent.core.jdbc.DruidPoolAgent.begin(jdbcUrl);" +
                            "return " + newBeforeExecuteMethodName + "();" +
                            "}");
            ctClass.addMethod(beforeExecuteMethodCopyMethod);

            CtMethod afterExecuteMethod = ctClass.getDeclaredMethod("afterExecute");
            CtMethod afterExecuteMethodCopyMethod = CtNewMethod.copy(afterExecuteMethod, ctClass, new ClassMap());
            String newAfterExecuteMethodName = afterExecuteMethodCopyMethod.getName() + "$Agent";
            afterExecuteMethod.setName(newAfterExecuteMethodName);
            afterExecuteMethodCopyMethod.setBody(
                    "{" +
                            "try {" +
                            "org.idea.qiyu.framework.apm.agent.core.jdbc.DruidPoolAgent.end();" +
                            "return " + newAfterExecuteMethodName + "();" +
                            "}" +
                            "finally {" +
                            "org.idea.qiyu.framework.apm.agent.core.jdbc.JdbcTraceSession.close();" +
                            "}" +
                            "}");
            ctClass.addMethod(afterExecuteMethodCopyMethod);
            return ctClass.toBytecode();
        } catch (Exception e) {
            e.printStackTrace();
        }
        return null;
    }

    /**
     * 这个步骤优先于begin和end函数
     *
     * @param sql
     */
    public static void recordSql(String sql) {
        JdbcStatisticsObject jdbcStatisticsObject = JdbcTraceSession.get();
        if (jdbcStatisticsObject == null) {
            jdbcStatisticsObject = new JdbcStatisticsObject();
            jdbcStatisticsObject.setQuerySql(sql);
            JdbcTraceSession.set(jdbcStatisticsObject);
        }
    }

    public static void begin(String jdbcUrl) {
        JdbcStatisticsObject jdbcStatisticsObject = JdbcTraceSession.get();
        if (jdbcStatisticsObject != null) {
            jdbcStatisticsObject.setBegin(System.currentTimeMillis());
            jdbcStatisticsObject.setJdbcUrl(jdbcUrl);
            jdbcStatisticsObject.setTraceId(UUID.randomUUID().toString());
        }
    }


    public static void end() {
        JdbcStatisticsObject jdbcStatisticsObject = JdbcTraceSession.get();
        if (jdbcStatisticsObject != null) {
            jdbcStatisticsObject.setEnd(System.currentTimeMillis());
            //sql查询耗时
            jdbcStatisticsObject.setCount(jdbcStatisticsObject.getEnd() - jdbcStatisticsObject.getBegin());
            System.out.println(jdbcStatisticsObject.toString());
        }
    }
}
